
# 导入tushare
import tushare as ts
import os
from dotenv import load_dotenv
from sqlalchemy import create_engine
import pymysql
import sqlalchemy
from loguru import logger

# 初始化pro接口
def main():
    # 读取.env文件
    load_dotenv(dotenv_path='.env')
    MYSQL_HOST = os.getenv('MYSQL_HOST')
    MYSQL_PORT = os.getenv('MYSQL_PORT')
    MYSQL_DATABASE = os.getenv('MYSQL_DATABASE')
    MYSQL_USER = os.getenv('MYSQL_USER')
    MYSQL_PASSWORD = os.getenv('MYSQL_PASSWORD')
    TUSHARE_API_KEY = os.getenv('TUSHARE_API_KEY')
    TRADE_DATE = os.getenv('TRADE_DATE', '')
    DAYS_TO_GET = int(os.getenv('DAYS_TO_GET', '10'))

    # 构建数据库连接
    engine = create_engine(
        f"mysql+pymysql://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_HOST}:{MYSQL_PORT}/{MYSQL_DATABASE}?charset=utf8mb4"
    )

    pro = ts.pro_api(TUSHARE_API_KEY)
    trade_date = TRADE_DATE
    days_to_get = DAYS_TO_GET

    for i in range(days_to_get):
        df = pro.daily(**{
            "ts_code": "",
            "trade_date": "",
            "start_date": "",
            "end_date": trade_date,
            "limit": "",
            "offset": ""
        }, fields=[
            "ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "change", "pct_chg", "vol", "amount"
        ])
        if df.empty:
            break
        # 字段重命名
        df = df.rename(columns={"change": "vchange"})
        trade_date = df.iloc[-1]['trade_date']
        current_trade_date = df.iloc[0]['trade_date']
        df=df[df['trade_date']==current_trade_date]
        # 保存到数据库
        try:
            df.to_sql('stock_daily', con=engine, if_exists='append', index=False)
            logger.info(f"Successfully saved data: {current_trade_date}")
        except pymysql.err.IntegrityError as e:
            logger.error(f"Error saving data: {current_trade_date}")
        except sqlalchemy.exc.IntegrityError as e:
            logger.error(f"Error saving data: {current_trade_date}")

if __name__ == "__main__":
    main()

        